#include <iostream>
#include <unordered_map>
#include <queue>
#include <fstream>

using namespace std;

struct HuffmanNode {
    char ch;
    int freq;
    HuffmanNode *left;
    HuffmanNode *right;
    HuffmanNode(char ch, int freq) : ch{ch}, freq{freq}, left{nullptr}, right{nullptr} {}

    bool operator>(const HuffmanNode& other) const {
        return freq > other.freq;
    }
};

unordered_map<char,int> StringCount(string str){
    unordered_map<char,int> count;
    for(auto c:str){
        count[c]++;
    }
    return count;
}

HuffmanNode* BuildHuffmanTree(unordered_map<char,int> freqMap) {
    priority_queue<HuffmanNode, vector<HuffmanNode>, greater<HuffmanNode>> minHeap;
    for(const auto &freq : freqMap) {
        minHeap.push(HuffmanNode{freq.first, freq.second});
    }
    while(minHeap.size() > 1) {
        HuffmanNode* left = new HuffmanNode(minHeap.top());
        minHeap.pop();
        HuffmanNode* right = new HuffmanNode(minHeap.top());
        minHeap.pop();

        HuffmanNode* merged = new HuffmanNode('\0', left->freq + right->freq);
        merged->left = left;
        merged->right = right;
        minHeap.push(*merged);
    }
    return new HuffmanNode(minHeap.top());
}

void GenerateHuffmanCode(HuffmanNode* root,const string &prefix, unordered_map<char, string> &codeMap){
    if(root == nullptr) return;
    if(root->ch != '\0'){ // 叶节点存储编码映射结果
        codeMap[root->ch] = prefix;
    }
    GenerateHuffmanCode(root->left, prefix + '0', codeMap);
    GenerateHuffmanCode(root->right, prefix + '1', codeMap);
}

string EncodeString(string str, unordered_map<char, string> &codeMap) {
    string encodedStr = "";
    for(const auto &c : str) {
        encodedStr += codeMap[c];
    }
    return encodedStr;
}

string DecodeString(string str, HuffmanNode* root) {
    string decodedStr = "";
    HuffmanNode *ptr = root;
    for(char &bit : str) {
        if(bit == '0') {
            ptr = ptr->left;
        } else {
            ptr = ptr->right;
        }
        if(ptr->ch != '\0') {
            decodedStr += ptr->ch;
            ptr = root;
        }
    }
    return decodedStr;
}

void TraversalTree(HuffmanNode *root, unordered_map<char, string> &codeMap){
    queue<HuffmanNode*> que;
    que.push(root);
    int i = 0;
    cout << endl;
    cout << "ptr" << "\t\t" << "index" << "\t" << "code" << "\t" << "freq" << "\t" << "char" << "\t" << "left" << "\t\t" << "right" << "\t\t" << endl;

    while(!que.empty()) {
        HuffmanNode *node = que.front();
        if(node == root){
            cout << "root" << "\t\t" << i << "\t" << codeMap[node->ch] << "\t" << node->freq << "\t" << node->ch << "\t";
        }else{
            cout << node << "\t" << i << "\t" << codeMap[node->ch] << "\t" << node->freq << "\t" << node->ch << "\t";
        }
        if(node->left == nullptr) {
            cout << "null";
            cout << "\t\t";
        } else {
            cout << node->left;
            cout << "\t";
        }
        if(node->right == nullptr) {
            cout << "null";
        } else {
            cout << node->right;
        }
        cout << endl;

        que.pop();
        if(node->left != nullptr) {
            que.push(node->left);
        }
        if(node->right != nullptr) {
            que.push(node->right);
        }
        i++;
    }
} 

int main(int argc, char const *argv[])
{
    cout << "the string convert to huffman:" << endl;
    string str{};
    getline(cin,str);
    unordered_map<char, int> count = StringCount(str);
    HuffmanNode* root = BuildHuffmanTree(count);
    unordered_map<char, string> codeMap;
    GenerateHuffmanCode(root, "", codeMap);
    string encodedStr = EncodeString(str, codeMap);
    cout << "the string convert to huffman code:" << encodedStr << endl;
    string decodedStr = DecodeString(encodedStr, root);
    cout << "the huffman code convert to string:" << decodedStr << endl;
    ofstream out("out.txt");
    if(out.is_open()){
        out << encodedStr;
        out.close();
    }else{
        cout << "file open error!" << endl;
    }
    TraversalTree(root, codeMap);
    for(auto &c : codeMap){
        if(c.first == '\0') continue;
        cout << c.first << ": " << c.second << "  ";
    }
    cout << endl;
    cout << encodedStr << endl;
    cout << decodedStr << endl;
    delete root;
    return 0;
}
